package main

import (
	"flag"
	"fmt"
	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/logger"
	"github.com/rs/zerolog/log"
	"gorm.io/gorm"
	"path/filepath"
	"shyxy-net/common/cache/lru"
	"shyxy-net/common/consul"
	commonModel "shyxy-net/common/db/model-service/model"
	"shyxy-net/model-service/config"
	"shyxy-net/model-service/service"
	result "shyxy-net/pkg/fiber"
)

var configDir = flag.String("config_dir", "./config", "config dir: default value is ./config")
var configFile = flag.String("config_file", "config.json", "config file: default value is config.json")

func main() {
	flag.Parse()

	configPath := filepath.Join(*configDir, *configFile)
	cfg, err := config.Loading(configPath)
	if err != nil {
		panic(err)
	}
	var port int
	port = cfg.ApplicationConfig.Port
	// 初始化数据库连接
	var modelInfoDB *gorm.DB
	var rpcServerDB *gorm.DB
	for _, mysql := range cfg.MysqlConfig {
		db, dbInitErr := mysql.InitDB()
		if dbInitErr != nil {
			panic(dbInitErr)
		}
		if mysql.Datasource.TableName == commonModel.ModelInfoTableName {
			modelInfoDB = db
		} else if mysql.Datasource.TableName == commonModel.RpcServerInfoTableName {
			rpcServerDB = db
		}
	}

	// 初始化模型服务
	modelService := service.New(modelInfoDB, rpcServerDB, lru.New(1000))

	// 启动服务注册
	address, serviceConfig, err := cfg.ConsulConfig.IsService()
	if err != nil {
		panic(err)
	}

	serviceRegister := consul.NewRegister(address, serviceConfig)
	if err := serviceRegister.Register(); err != nil {
		log.Error().Msgf("Failed to register service: %v", err)
	}
	defer func() {
		if err := serviceRegister.Deregister(); err != nil {
			log.Error().Msgf("Failed to deregister service: %v", err)
		}
	}()

	// 封装模型服务路由
	app := fiber.New()
	app.Use(logger.New())
	setRoutes(app, modelService)

	// 启动 HTTP 服务
	name := cfg.ApplicationConfig.Name
	log.Printf("Starting %s on %d", name, port)
	if err := app.Listen(fmt.Sprintf(":%d", port)); err != nil {
		log.Error().Msgf("Failed to start %s HTTP server: %v", name, err)
	}
}

func setRoutes(app *fiber.App, modelService *service.ModelService) {
	app.Get("/health", func(c *fiber.Ctx) error {
		return c.SendString("OK")
	})
	app.Get("/models", func(c *fiber.Ctx) error {
		models, err := modelService.ListModels()
		if err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(result.FailedWithErrCode(result.NetworkError))
		}
		return c.Status(fiber.StatusOK).JSON(result.Succeed(fiber.Map{
			"models": models,
		}))
	})
	app.Get("/models/:modelName/exists", func(c *fiber.Ctx) error {
		modelName := c.Params("modelName")
		exists, err := modelService.CheckIfModelExists(modelName)
		log.Info().Msgf("modelName: %s, exists: %v", modelName, exists)
		if err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(result.FailedWithErrCode(result.NetworkError))
		}
		return c.Status(fiber.StatusOK).JSON(result.Succeed(fiber.Map{
			"exists": exists,
		}))
	})
	app.Get("/models/:modelName/rpc", func(c *fiber.Ctx) error {
		modelName := c.Params("modelName")
		addr, err := modelService.GetRpcServerAddressByModelName(modelName)
		log.Info().Msgf("modelName: %s, backend: %+v", modelName, addr)
		if err != nil {
			return c.Status(fiber.StatusInternalServerError).JSON(result.FailedWithErrCode(result.NetworkError))
		}
		return c.Status(fiber.StatusOK).JSON(result.Succeed(fiber.Map{
			"address": addr,
		}))
	})
}
